package clientcache

import (
    "cing.io/cing/internal/config"
    "cing.io/cing/internal/k8s"
    "cing.io/cing/internal/storage"
    "cing.io/cing/pkg/util"
    "context"
    "fmt"
    log "github.com/sirupsen/logrus"
    "google.golang.org/grpc"
    "google.golang.org/grpc/connectivity"
    "google.golang.org/grpc/credentials/insecure"
    v1 "k8s.io/api/core/v1"
    "k8s.io/apimachinery/pkg/util/sets"
    "sync"
    "time"

    "cing.io/cing-api/singer"
)

var Singer *singerClientCache

type singerClientCache struct {
    holderMap map[string]*singerHolder // key: svc full name, val: holder
    m         sync.Mutex               // mutex for key setting
}

type singerHolder struct {
    m       sync.RWMutex
    cliSet  map[string]singer.CingSingerClient // key: ip, val: cli
    connSet map[string]*grpc.ClientConn        // key: ip, val: conn

    closeFuncSet          map[string]map[string]func()  // key: ip, val: map(k:timestamp/taskFullName,v:close function list)
    closeCtxCancelFuncSet map[string]context.CancelFunc // key: ip, val: ctx cancel function
}

func (cache *singerClientCache) GetSingerSet(taskNamespace, taskName string) map[string]singer.CingSingerClient {
    task := storage.GetTask(taskNamespace, taskName)
    sns, svc := util.GetTaskSingerService(task)
    holder, err := cache.getHolderInService(sns, svc)
    if err != nil {
        log.Error("get planner client err:", err)
        return nil
    }
    return holder.cliSet
}

func (cache *singerClientCache) Update(endpoints *v1.Endpoints) {
    SvcFullName := util.MergeNamespaceAndName(endpoints.Namespace, endpoints.Name)
    holder, exists := cache.holderMap[SvcFullName]
    if !exists || holder == nil {
        return
    }

    holder.m.Lock()
    defer holder.m.Unlock()

    newIpList := util.GetIpListFromEndpoints(endpoints)

    newIpSet := sets.NewString(newIpList...)

    // close conn that does not exist in the new ip list
    for ip := range holder.connSet {
        _, exists := newIpSet[ip]
        if !exists {
            if err := holder.clear(ip); err != nil {
                log.Error("grpc conn close error: ", err)
            }
        }
    }

    // open conn that exist in the new ip list but does not exist in conn set
    for ip := range newIpSet {
        _, exists := holder.connSet[ip]
        if !exists {
            conn, err := grpc.Dial(
                fmt.Sprintf("%s:%d", ip, config.SingerPort),
                grpc.WithTransportCredentials(insecure.NewCredentials()),
            )
            if err != nil {
                log.Error("grpc connect error: ", err)
                continue
            }
            holder.connSet[ip] = conn
            holder.cliSet[ip] = singer.NewCingSingerClient(conn)
        }
    }
}

// Clear to clear the cinger client set
func (cache *singerClientCache) Clear() {
    for _, holder := range cache.holderMap {
        for ip := range holder.connSet {
            if err := holder.clear(ip); err != nil {
                log.Error("grpc conn close error: ", err)
            }
        }
    }
    cache.holderMap = make(map[string]*singerHolder)
}

func (cache *singerClientCache) getHolderInService(namespace, service string) (*singerHolder, error) {
    svcFullName := util.MergeNamespaceAndName(namespace, service)
    holder := cache.holderMap[svcFullName]
    if holder == nil {
        cache.m.Lock()
        if holder = cache.holderMap[svcFullName]; holder == nil {
            holder = &singerHolder{}
            cache.holderMap[svcFullName] = holder
        }
        cache.m.Unlock()
    }
    holder.m.RLock()
    cliSet := holder.cliSet
    if cliSet == nil {
        holder.m.RUnlock()
        holder.m.Lock()
        if cliSet = holder.cliSet; cliSet == nil {
            err := holder.initSingerSet(namespace, service)
            if err != nil {
                return nil, err
            }
        }
        holder.m.Unlock()
        return holder, nil
    }
    holder.m.RUnlock()
    return holder, nil
}

func (cache *singerClientCache) AddCloseFunc(timestamp int64, taskNamespace, taskName, ip string, fun func()) {
    task := storage.GetTask(taskNamespace, taskName)
    svcFullName := util.MergeNamespaceAndName(util.GetTaskSingerService(task))
    holder := cache.holderMap[svcFullName]
    holder.m.Lock()
    defer holder.m.Unlock()

    key := fmt.Sprintf("%d/%s/%s", timestamp, taskNamespace, taskName)
    holder.closeFuncSet[ip][key] = fun
}

func (cache *singerClientCache) DeleteCloseFunc(timestamp int64, taskNamespace, taskName, ip string) {
    task := storage.GetTask(taskNamespace, taskName)
    svcFullName := util.MergeNamespaceAndName(util.GetTaskSingerService(task))
    holder := cache.holderMap[svcFullName]
    holder.m.Lock()
    defer holder.m.Unlock()

    key := fmt.Sprintf("%d/%s/%s", timestamp, taskNamespace, taskName)
    holder.closeFuncSet[ip][key] = nil
}

func (holder *singerHolder) initSingerSet(namespace, service string) error {
    if config.LocalTest {
        return holder.localTestInitSinger()
    } else {
        return holder.inClusterInitSinger(namespace, service)
    }
}

// the only way of close grpc conn
func (holder *singerHolder) clear(ip string) error {
    conn := holder.connSet[ip]
    state := conn.GetState()
    if state == connectivity.TransientFailure || state == connectivity.Shutdown {
        if err := conn.Close(); err != nil {
            log.Error("grpc conn close error: ", err)
        }
        holder.cliSet[ip] = nil
        holder.connSet[ip] = nil

        // cancel WaitForStateChange goroutine
        holder.closeCtxCancelFuncSet[ip]()
    }
    return nil
}

func (holder *singerHolder) inClusterInitSinger(namespace, service string) error {
    endpoints, err := k8s.EndpointsInformer().Lister().Endpoints(namespace).Get(service)
    if err != nil {
        return err
    }

    ipList := util.GetIpListFromEndpoints(endpoints)

    holder.cliSet = make(map[string]singer.CingSingerClient)
    holder.connSet = make(map[string]*grpc.ClientConn)

    holder.closeFuncSet = make(map[string]map[string]func())
    holder.closeCtxCancelFuncSet = make(map[string]context.CancelFunc)

    for _, ip := range ipList {
        conn, err := grpc.Dial(
            fmt.Sprintf("%s:%d", ip, config.SingerPort),
            grpc.WithTransportCredentials(insecure.NewCredentials()),
        )
        if err != nil {
            return err
        }

        holder.connSet[ip] = conn
        holder.cliSet[ip] = singer.NewCingSingerClient(conn)

        holder.closeFuncSet[ip] = make(map[string]func())
        ctx, cancel := context.WithCancel(context.Background())
        holder.closeCtxCancelFuncSet[ip] = cancel

        // handle grpc connection state change
        go watchConnectionState(ctx, holder.closeFuncSet[ip], conn)
    }

    return nil
}

func watchConnectionState(ctx context.Context, set map[string]func(), conn *grpc.ClientConn) {
    for {
        select {
        case <-ctx.Done():
            return
        default:
            if conn.WaitForStateChange(ctx, conn.GetState()) {
                if conn.GetState() == connectivity.Shutdown {
                    for _, fun := range set {
                        fun()
                    }
                } else if connNotAlive(conn) {
                    time.Sleep(config.SingerGrpcTimeout)
                    if connNotAlive(conn) {
                        for _, fun := range set {
                            fun()
                        }
                    }
                }
            }
        }
    }
}

func connNotAlive(conn *grpc.ClientConn) bool {
    state := conn.GetState()
    return state == connectivity.TransientFailure
}

func (holder *singerHolder) localTestInitSinger() error {
    conn, err := grpc.Dial(
        fmt.Sprintf("%s:%d", config.Localhost, config.SingerPort),
        grpc.WithTransportCredentials(insecure.NewCredentials()),
    )
    if err != nil {
        return err
    }

    cli := singer.NewCingSingerClient(conn)

    holder.connSet = map[string]*grpc.ClientConn{config.Localhost: conn}
    holder.cliSet = map[string]singer.CingSingerClient{config.Localhost: cli}

    return nil
}
